import threading

import cv2
import mss
import numpy as np


class ScreenCapture:
    def __init__(self):
        self.thread_local = threading.local()

    def _get_mss_instance(self) -> mss:
        """获取当前线程的 mss 实例"""
        if not hasattr(self.thread_local, 'mss'):
            self.thread_local.mss = mss.mss()
        return self.thread_local.mss

    def screen_shot(self, screen_shot_area=(0, 0, 1067, 600)):
        x, y, width, height = screen_shot_area
        monitor = {"top": y, "left": x, "width": width, "height": height}
        sct = self._get_mss_instance()
        sct_img = np.array(sct.grab(monitor))
        if sct_img.shape[2] == 4:  # 如果有Alpha通道
            sct_img = cv2.cvtColor(sct_img, cv2.COLOR_BGRA2BGR)
        return sct_img

    def screen_shot_all(self, screen_shot_area=(0, 0, 1067, 600)):
        x, y, width, height = screen_shot_area
        monitor = {"top": y, "left": x, "width": width, "height": height}
        sct = self._get_mss_instance()
        img = sct.grab(monitor)
        sct_img = np.array(img)
        if sct_img.shape[2] == 4:  # 如果有Alpha通道
            sct_img = cv2.cvtColor(sct_img, cv2.COLOR_BGRA2BGR)
        img_gray = cv2.cvtColor(sct_img, cv2.COLOR_BGR2GRAY)
        return sct_img, img_gray

    def release(self):
        self.sct.close()